import torch 
from PIL import Image

import cn_clip.clip as clip
from cn_clip.clip import load_from_name, available_models



# 打印可用的模型名称
print("Available models:", available_models(),"\n")  
# 输出: Available models: ['ViT-B-16', 'ViT-L-14', 'ViT-L-14-336', 'ViT-H-14', 'RN50']

# 判断是否有可用的CUDA设备（GPU），否则使用CPU
device = "cuda" if torch.cuda.is_available() else "cpu"

# 从预训练模型中加载指定名称的模型（这里选择的是RN50），并设置下载根目录
model, preprocess = load_from_name("RN50", device=device, download_root='./')

# 设置模型为评估模式
model.eval()


# 然后将处理后的图像转换为批次形式（unsqueeze(0)），并移动到选定的计算设备（GPU或CPU）
image = preprocess(Image.open("6.jpg")).unsqueeze(0).to(device)

# 使用CLIP的tokenize方法将文本转化为模型可接受的格式，并移动到选定的计算设备
text = clip.tokenize(["按钮","圆孔"]).to(device)

# 在不进行梯度计算的情况下执行以下操作以节省内存和计算资源
with torch.no_grad():
    # 编码图像和文本，得到各自的特征向量
    image_features = model.encode_image(image)
    text_features = model.encode_text(text)
    
    # 对提取出的图像和文本特征进行归一化处理
    image_features /= image_features.norm(dim=-1, keepdim=True) 
    text_features /= text_features.norm(dim=-1, keepdim=True)    

    # logits_per_image 表示图像与文本之间的相似度分数。logits_per_text 表示文本与图像之间的相似度分数。
    # logits_per_image 和 logits_per_text 的值相同是因为它们表示的是相同的相似度信息，只是从不同的视角（图像对文本 vs 文本对图像）来表示。
    logits_per_image, logits_per_text = model.get_similarity(image, text)
    
    # 将相似度分数通过softmax函数转换为概率分布，并将结果移到CPU上并转换为numpy数组
    probs = logits_per_image.softmax(dim=-1).cpu().numpy()

# 打印每个标签的概率
print("Label probs:", probs)
print(logits_per_image, logits_per_text)
